import sys
import json
import os

from getHWconfig import get_HW_configuration
from generate_md import generate_markdown
from generate_compare import generate_blas_performance_chart

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))

def load_benchmark_data(file_path):
    with open(file_path) as f:
        data = json.load(f)
    return data


def create_module_dict(data):
    module_dict = {}
    for item in data['result']:
        module = item.get('module')
        if module:
            module_dict[module] = {
                'value': item.get('metric_value', 'N/A'),
                'unit': item.get('metric_unit', 'N/A')
            }
    return module_dict

def process_data(dict_pair_1, dict_pair_2):
    dict1 = dict_pair_1[1]
    dict2 = dict_pair_2[1]
    dict1_name = dict_pair_1[0]
    dict2_name = dict_pair_2[0]

    table_data = []

    for module in set(dict1) & set(dict2):
        row = [
            module,
            dict1[module]['value'],
            dict1[module]['unit'],
            dict2[module]['value'],
            dict2[module]['unit']
        ]
        table_data.append(row)

    table_data.sort(key=lambda x: x[0])

    header = ['Module', dict1_name, 'Unit', dict2_name, 'Unit', 'pct']
    rows = []

    for module in sorted(set(dict1.keys()) & set(dict2.keys())):
        if (dict1[module]['unit'] == 'us' or dict1[module]['unit'] == 'ms' or dict1[module]['unit'] == 'clk') and not float(dict2[module]['value']) < 0.1:
            pecent = float(dict1[module]['value']) / float(dict2[module]['value']) * 100
        elif not float(dict1[module]['value']) < 0.1:
            pecent = float(dict2[module]['value']) / float(dict1[module]['value']) * 100
        else:
            pecet = 0
        row = [
            module,
            dict1[module]['value'],
            dict1[module]['unit'],
            dict2[module]['value'],
            dict2[module]['unit'],
            str(int(pecent * 1000)/1000) + '%'
        ]
        rows.append(row)

    return [header] + rows


def mixed_unit(data, is_single=False):
    if not is_single:
        return [[data[0][0], data[0][1], data[0][3], data[0][5]]] + [[v[0], str(v[1]) + ' ' + str(v[2]), str(v[3]) + ' ' + str(v[4]), v[5]] for v in data[1:]]
    else:
        return [[data[0][0], data[0][1]]] + [[v[0], str(v[1]) + ' ' + str(v[2])] for v in data[1:]]


def gen_single_table(data):
    header = ['Module', 'Performance']
    rows = []
    for module in sorted(set(data.keys())):
        row = [module, data[module]['value'], data[module]['unit']]
        rows.append(row)
    return [header] + rows


def gen_comparative_table():
    pass


def main():
    if len(sys.argv) == 1 or len(sys.argv) > 3:
        print("""Usage: python script.py <first_json_file> <second_json_file> to generate
       python script.py <single_json_file>""")
        sys.exit(1)

    hw_info = get_HW_configuration()

    if len(sys.argv) == 2:
        data_path = sys.argv[1]
        data_table = gen_single_table(create_module_dict(load_benchmark_data(data_path)))
        generate_markdown(hw_info, mixed_unit(data_table, True), f"{SCRIPT_DIR}/../report/", False)
        return

    nv_json_file_path = sys.argv[1]
    mx_json_file_path = sys.argv[2]

    mx_json_data = load_benchmark_data(mx_json_file_path)
    nv_json_data = load_benchmark_data(nv_json_file_path)

    mx_data = create_module_dict(mx_json_data)
    nv_data = create_module_dict(nv_json_data)

    compare_result = process_data([nv_json_data['gpu_name'], nv_data], [mx_json_data['gpu_name'], mx_data])

    generate_blas_ok = generate_blas_performance_chart(compare_result, f"{SCRIPT_DIR}/../report/blas_performance_comparison.png", nv_json_data['gpu_name'], mx_json_data['gpu_name'])
    generate_markdown(hw_info, mixed_unit(compare_result), f"{SCRIPT_DIR}/../report/", generate_blas_ok)


if __name__ == "__main__":
    main()
